---
title: "Estimate reproduction number with only conversational contacts included (sensitivity analysis)"
output:
  html_document:
    df_print: paged
---

```{r setup, include=FALSE}
library(socialmixr)
library(dplyr)
library(here)
library(ggplot2)
library(tidyverse)
library(ggthemes)
library(cowplot)
library(reshape2)
library(ipumsr)
library(grid)
library(extrafont)

# load some functions
source(here('bics-paper-code', 'code', 'utils.R'))

```

Load all data
```{r}
## survey data
# set polymod parameters for fetching data
polymod_country = "United Kingdom"
polymod_age_limits = c(0, 18, 25, 35, 45,  65)

# polymod
polymod_mat <- getPolymodMatrix(polymod_country = polymod_country,
                                polymod_age_limits= polymod_age_limits)

polymod_mat_fb_age <- getPolymodMatrix(polymod_country = polymod_country,
                                polymod_age_limits= c(0, 15, 25, 35, 45,  65))


# polymod without school contacts
polymod_no_school_mat <- getPolymodMatrixNoSchool(polymod_country = polymod_country,
                                                  polymod_age_limits= polymod_age_limits)


# grab age distns with kids included for making contact matrix symmetric
acs18_agecat_withkids_targets <- readRDS(file=here('bics-paper-code', 'data', 'ACS', 'acs18_wave1_agecat_withkids.rds'))

# grab 2015 acs data for the FB matrix
acs15_agecat_fb <- readRDS(file=here('bics-paper-code', 'data', 'ACS', 'acs15_fb_agecat_withkids.rds'))

## wave 3 data 
wave3 <- readRDS(file=here('bics-paper-code', 'data', 'contact-matrices', 'wave3_contact_matrix_w0age_onlycc.rds')) %>%
  rename(ego_age =.ego_age, alter_age=.alter_age) %>%
  filter(ego_age != "[0,18)") %>%
  filter(!is.na(alter_age))


## wave 2 data 
wave2 <- readRDS(file=here('bics-paper-code', 'data', 'contact-matrices', 'wave2_contact_matrix_w0age_onlycc.rds')) %>%
  rename(ego_age =.ego_age, alter_age=.alter_age) %>%
  filter(ego_age != "[0,18)") %>%
  filter(!is.na(alter_age))


## wave 1 data
wave1 <- readRDS(file=here('bics-paper-code', 'data', 'contact-matrices', 'wave1_contact_matrix_w0age_onlycc.rds')) %>%
  rename(ego_age =.ego_age, alter_age=.alter_age) %>%
  filter(ego_age != "[0,18)") 

## wave 0 data
wave0 <- readRDS(here('bics-paper-code', 'data', 'contact-matrices', "wave0_contact_matrix_w0age.rds")) %>%
  rename(ego_age =.ego_age, alter_age=.alter_age) %>%
  filter(ego_age != "[0,18)") %>%
  filter(!is.na(alter_age))

## create a list object with the data from each wave
wave_df_list <- list(wave0, wave1, wave2, wave3)
num_waves <-length(wave_df_list)

## 2015 facebook study
fb <- readRDS(file=here('bics-paper-code', 'data', 'contact-matrices', 'fb_contact_matrix_fbage.rds')) %>%
  rename(ego_age = .ego_age, alter_age = .alter_age)  %>%
  filter(ego_age != "[0,15)") 
```

Adjust for reciprocity and infer within group contacts for youngest age group

```{r}
# BICS
filled_matrix_BICS <- survey_mat_youngest_added_BICS <- list()
for (n in 1:num_waves) {
  filled_matrix_BICS[[n]] <- fill_matrix_BICS(df = wave_df_list[[n]], 
                                              age_df = acs18_agecat_withkids_targets,
                                              fill_mat = polymod_no_school_mat)
  survey_mat_youngest_added_BICS[[n]] <- filled_matrix_BICS[[n]]$survey_mat_youngest_added
}

# 2015 FB study (baseline)
filled_matrix_fb <- fill_matrix_FB(df = fb, 
                                     age_df = acs15_agecat_fb,
                                     fill_mat = polymod_mat_fb_age) # fill FB matrix assuming kids are going to school (i.e. business as usual)
survey_mat_youngest_added_fb<- filled_matrix_fb$survey_mat_youngest_added


```

Relative ratios of eigenvalues 

```{r}
#BICS compared to polymod
ratio_BICS_polymod <- list()
for (n in 1:num_waves){
  ratio_BICS_polymod[[n]] <- getRelativeR0(survey_mat = survey_mat_youngest_added_BICS[[n]],
                                           comparison_mat = polymod_mat)
}

#BICS compared to 2015 FB study
ratio_BICS_fb <- list()
for (n in 1:num_waves){
  ratio_BICS_fb[[n]] <- getRelativeR0(survey_mat = survey_mat_youngest_added_BICS[[n]],
                                           comparison_mat = survey_mat_youngest_added_fb)
}

```

Load bootstrap data and calculate \(R_0\) and confidence intervals

```{r}
BICS_bootstrapped_wave0 <- readRDS(file=here('bics-paper-code', 'data', 'contact-matrices','wave0_contact_matrices_w0age_bootstrapped.rds'))
BICS_bootstrapped_wave1 <- readRDS(file=here('bics-paper-code', 'data', 'contact-matrices', 'wave1_contact_matrices_w0age_onlycc_bootstrapped.rds'))
BICS_bootstrapped_wave2 <- readRDS(file=here('bics-paper-code', 'data', 'contact-matrices', 'wave2_contact_matrices_w0age_onlycc_bootstrapped.rds'))
BICS_bootstrapped_wave3 <- readRDS(file=here('bics-paper-code', 'data', 'contact-matrices', 'wave3_contact_matrices_w0age_onlycc_bootstrapped.rds'))

wave_bootstrapped_df_list <- list(BICS_bootstrapped_wave0,BICS_bootstrapped_wave1,BICS_bootstrapped_wave2,
                                  BICS_bootstrapped_wave3)

polymod_no_school_mat_bootstrapped <- readRDS(file=here('bics-paper-code', 'data', 'polymod', 'polymod_no_school_bootstrapped.rds'))
polymod_mat_bootstrapped <- readRDS(file=here('bics-paper-code', 'data', 'polymod', 'polymod_bootstrapped.rds'))
polymod_mat_fb_age_bootstrapped <- readRDS(file=here('bics-paper-code', 'data', 'polymod', 'polymod_mat_fb_age_bootstrapped.rds'))

fb_bootstrapped <- readRDS(file=here('bics-paper-code', 'data', 'contact-matrices','fb_contact_matrices_fbage_bootstrapped.rds'))


baseline_R0_bootstrapped <- readRDS(file = here('bics-paper-code', 'data', 'polymod', 'baseline_r_bootstrapped.rds'))


baseline_R0 = 2.5 
baseline_R0_SE = 0.54
num_replicates <- length(fb_bootstrapped)

ratio_bootstrapped_polymod_baseline <- R0_bootstrapped_polymod_baseline <- matrix(NA, ncol = num_waves, nrow = num_replicates)
ratio_bootstrapped_fb_baseline <- R0_bootstrapped_fb_baseline <- matrix(NA, ncol = num_waves, nrow = num_replicates)

# for average contact by age barplot ci
sum_sym_avg_fb_bootstrapped <- list()
sum_sym_avg_BICS_bootstrapped <- list()

for (i in 1:num_replicates) {
  
  
  # infer mixing in youngest age group
  filled_matrix_fb <- fill_matrix_FB(df = fb_bootstrapped[[i]] %>% ungroup() %>%
                                       rename(ego_age = .ego_age, alter_age = .alter_age) %>% 
                                       filter(ego_age != "[0,15)"), 
                                     age_df = acs15_agecat_fb,
                                     fill_mat = polymod_mat_fb_age_bootstrapped[[i]]$matrix)
  
  survey_mat_youngest_added_fb_tmp <- filled_matrix_fb$survey_mat_youngest_added 
  sum_sym_avg_fb_bootstrapped[[i]] <- rowSums(survey_mat_youngest_added_fb_tmp)
  
  # draw baseline R0 value
  baselineR0_tmp <- baseline_R0_bootstrapped[[i]]  
  
  
  sum_sym_avg_BICS_bootstrapped_tmp <- list()
  
  for (n in 1:num_waves) {
    filled_matrix <- fill_matrix_BICS(df = wave_bootstrapped_df_list[[n]][[i]] %>% ungroup() %>%
                                        mutate(ego_age = .ego_age, alter_age = .alter_age) %>% 
                                        filter(ego_age != "[0,18)") %>%
                                        filter(!is.na(alter_age)), 
                                      age_df = acs18_agecat_withkids_targets,
                                      fill_mat = polymod_no_school_mat_bootstrapped[[i]]$matrix)
    survey_mat_youngest_added_tmp <- filled_matrix$survey_mat_youngest_added
    
    sum_sym_avg_BICS_bootstrapped_tmp[[n]] <- rowSums(survey_mat_youngest_added_tmp)
    
    
    # polymod as baseline
    ratio_bootstrapped_polymod_baseline[i,n] = getRelativeR0(survey_mat = survey_mat_youngest_added_tmp,
                                                             comparison_mat =polymod_mat_bootstrapped[[i]]$matrix)
    
    R0_bootstrapped_polymod_baseline[i,n] <- ratio_bootstrapped_polymod_baseline[i,n]*baselineR0_tmp
    
    # fb as baseline
    ratio_bootstrapped_fb_baseline[i,n] = getRelativeR0(survey_mat = survey_mat_youngest_added_tmp,
                                                        comparison_mat = survey_mat_youngest_added_fb_tmp)
    
    R0_bootstrapped_fb_baseline[i,n] <- ratio_bootstrapped_fb_baseline[i,n]*baselineR0_tmp
    
    
    
  }
  
  sum_sym_avg_BICS_bootstrapped[[i]] <- sum_sym_avg_BICS_bootstrapped_tmp
  
}


```

\(R_0\) estimate and 95% confidence intervals

```{r}
# polymod baseline: R0 est + ci
R0_est_polymod_baseline <- lapply(ratio_BICS_polymod, function(x) x*baseline_R0)
ci_R0_est_polymod_baseline <- list()
for (n in 1:num_waves) {
  ci_R0_est_polymod_baseline[[n]] <- calc_R0_ci_percentile(R0_bootstrapped_est = R0_bootstrapped_polymod_baseline[,n])
  
}

#fb baseline: R0 est + ci
R0_est_fb_baseline <- lapply(ratio_BICS_fb, function(x) x*baseline_R0)
ci_R0_est_fb_baseline <- list()
for (n in 1:num_waves) {
  ci_R0_est_fb_baseline[[n]] <- calc_R0_ci_percentile(R0_bootstrapped_est = R0_bootstrapped_fb_baseline[,n])
  
}


# % decline
(1 - unlist(ratio_BICS_fb))*100

# % decline in R0 : CI
ci_ratio_fb_baseline <- list()
for (n in 1:num_waves){
  ci_ratio_fb_baseline[[n]] <- (1-calc_R0_ci_percentile(R0_bootstrapped_est = ratio_bootstrapped_fb_baseline[,n]))*100
}

```

Plot \(R_0\) estimates

```{r}
plot_df_r0_ci <- as.data.frame(cbind(type = "Baseline", comparison = "Baseline", 
                                     r0 = baseline_R0, 
                                     ci_low = baseline_R0 - 1.96*baseline_R0_SE, 
                                     ci_high = baseline_R0 + 1.96*baseline_R0_SE)) %>%
  rbind(cbind(type = "Wave 0", comparison = "FB",
              r0 = R0_est_fb_baseline[[1]], 
              ci_low = ci_R0_est_fb_baseline[[1]][1], 
              ci_high = ci_R0_est_fb_baseline[[1]][2])) %>%
  rbind(cbind(type = "Wave 1", comparison = "FB",
              r0 = R0_est_fb_baseline[[2]], 
              ci_low = ci_R0_est_fb_baseline[[2]][1], 
              ci_high = ci_R0_est_fb_baseline[[2]][2])) %>%
  rbind(cbind(type = "Wave 2", comparison = "FB",
              r0 = R0_est_fb_baseline[[3]], 
              ci_low = ci_R0_est_fb_baseline[[3]][1], 
              ci_high = ci_R0_est_fb_baseline[[3]][2])) %>%
  rbind(cbind(type = "Wave 3", comparison = "FB",
              r0 = R0_est_fb_baseline[[4]], 
              ci_low = ci_R0_est_fb_baseline[[4]][1], 
              ci_high = ci_R0_est_fb_baseline[[4]][2])) %>%
  rbind(cbind(type = "Wave 0", comparison = "POLYMOD",
              r0 = R0_est_polymod_baseline[[1]], 
              ci_low = ci_R0_est_polymod_baseline[[1]][1], 
              ci_high = ci_R0_est_polymod_baseline[[1]][2])) %>%
  rbind(cbind(type = "Wave 1", comparison = "POLYMOD",
              r0 = R0_est_polymod_baseline[[2]], 
              ci_low = ci_R0_est_polymod_baseline[[2]][1], 
              ci_high = ci_R0_est_polymod_baseline[[2]][2])) %>%
  rbind(cbind(type = "Wave 2", comparison = "POLYMOD",
              r0 = R0_est_polymod_baseline[[3]], 
              ci_low = ci_R0_est_polymod_baseline[[3]][1], 
              ci_high = ci_R0_est_polymod_baseline[[3]][2])) %>%
  rbind(cbind(type = "Wave 3", comparison = "POLYMOD",
              r0 = R0_est_polymod_baseline[[4]], 
              ci_low = ci_R0_est_polymod_baseline[[4]][1], 
              ci_high = ci_R0_est_polymod_baseline[[4]][2])) %>%
  mutate(r0 = as.numeric(as.character(r0)), ci_low = as.numeric(as.character(ci_low)), ci_high=as.numeric(as.character(ci_high)))

write.csv(plot_df_r0_ci, file = here("bics-paper-code", "out","figure_S5_R0_figure_onlycc.csv"), row.names = FALSE)


R0plot <- plot_df_r0_ci %>% 
  ggplot(aes(x=type, y = r0, group = comparison, color = comparison)) +
  scale_color_manual(name = "Relative to", values = c("black", "blue", "maroon")) +
  geom_errorbar(width=.1, aes(ymin=ci_low, ymax=ci_high), position=position_dodge(width=0.5)) +
  geom_point(shape=19, size=3, fill="white", position=position_dodge(width=0.5)) +
  geom_hline(yintercept = 1, col = "black", linetype = "dotted")+
  labs(y = expression(R[0]~estimate), x = "")+
  theme_classic(base_size = 8) +
  theme(legend.position="bottom")

ggsave(R0plot, file = here("bics-paper-code", "out", "R0_figure_onlycc.png"), width = 6, height = 4)
ggsave(R0plot, file = here("bics-paper-code", "out", "R0_figure_onlycc.pdf"), width = 6, height = 4)
ggsave(R0plot, file = here("bics-paper-code", "out", "figure_S5.pdf"), width = 6, height = 4)


R0plot


```

